#include "TopLevelAccelerationStructure.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

TopLevelAccelerationStructure::Builder::Builder() noexcept = default;
TopLevelAccelerationStructure::Builder::~Builder() noexcept = default;
TopLevelAccelerationStructure::Builder::Builder(TopLevelAccelerationStructure::Builder const& rhs) noexcept = default;
TopLevelAccelerationStructure::Builder::Builder(TopLevelAccelerationStructure::Builder && rhs) noexcept = default;
TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::operator=(TopLevelAccelerationStructure::Builder const& rhs) noexcept = default;
TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::operator=(TopLevelAccelerationStructure::Builder && rhs) noexcept = default;

TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::SetFlags(VkBuildAccelerationStructureFlagsKHR flags) noexcept
{
	mImpl->flags = flags;
	return *this;
}

TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::SetArrayOfPointers(VkBool32 arrayOfPointers) noexcept
{
	mImpl->arrayOfPointers = arrayOfPointers;
	return *this;
}

TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::SetGeometryFlags(VkGeometryFlagsKHR geometryFlags) noexcept
{
	mImpl->geometryFlags = geometryFlags;
	return *this;
}

TopLevelAccelerationStructure::Builder& TopLevelAccelerationStructure::Builder::AddBlas(RefCountPtr<BottomLevelAccelerationStructure> bottomLevelAccelerationStructure) noexcept
{
	mImpl->bottomLevelAccelerationStructures.push_back(bottomLevelAccelerationStructure);
	return *this;
}

RefCountPtr<TopLevelAccelerationStructure> TopLevelAccelerationStructure::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<TopLevelAccelerationStructure>(NewObject<TopLevelAccelerationStructure>(*mImpl));
}

RefCountPtr<TopLevelAccelerationStructure> TopLevelAccelerationStructure::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	TopLevelAccelerationStructure* topLevelAccelerationStructure = NewObject<TopLevelAccelerationStructure>(*mImpl);
	process->AddAccelerationStructure(name, topLevelAccelerationStructure);
	return RefCountPtr<TopLevelAccelerationStructure>(topLevelAccelerationStructure);
}

void TopLevelAccelerationStructure::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveAccelerationStructure(mName, this);
		}
		DeleteObject<TopLevelAccelerationStructure>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveAccelerationStructure(mName, this);
			}
			DeleteObject<TopLevelAccelerationStructure>(this);
			});
	}
}

TopLevelAccelerationStructure::TopLevelAccelerationStructure(BuilderDetails & builderDetails)
{
	mType = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR;
	mFlags = builderDetails.flags;

	mAccelerationStructureGeometry = {};
	mAccelerationStructureGeometry.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR;
	mAccelerationStructureGeometry.geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR;
	mAccelerationStructureGeometry.flags = builderDetails.geometryFlags;
	mAccelerationStructureGeometry.geometry.instances.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR;
	mAccelerationStructureGeometry.geometry.instances.arrayOfPointers = builderDetails.arrayOfPointers;

	mBottomLevelAccelerationStructures.assign(builderDetails.bottomLevelAccelerationStructures.begin(), builderDetails.bottomLevelAccelerationStructures.end());


	// make sure all the blas geometries are build and compaction
	CompactionAndBuildBlas();

	// collect the instance and craete instance buffer
	CreateInstanceGeometries();



	uint32 maxPrimitiveCount = static_cast<uint32>(mBottomLevelAccelerationStructures.size());

	VkAccelerationStructureBuildRangeInfoKHR accelerationStructureBuildRangeInfo{};
	accelerationStructureBuildRangeInfo.primitiveCount = maxPrimitiveCount;
	accelerationStructureBuildRangeInfo.primitiveOffset = 0;
	accelerationStructureBuildRangeInfo.firstVertex = 0;
	accelerationStructureBuildRangeInfo.transformOffset = 0;

	std::array<VkAccelerationStructureBuildRangeInfoKHR*, 1> accelerationStructureBuildRangeInfos = { &accelerationStructureBuildRangeInfo };

	// Get Size Info
	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureSizeGeometryInfo{};
	accelerationStructureSizeGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureSizeGeometryInfo.type = mType;
	accelerationStructureSizeGeometryInfo.flags = mFlags;
	// Tlas geometryCount must be 1
	accelerationStructureSizeGeometryInfo.geometryCount = 1;
	accelerationStructureSizeGeometryInfo.pGeometries = &mAccelerationStructureGeometry;

	VkAccelerationStructureBuildSizesInfoKHR accelerationStructureBuildSizesInfo{};
	accelerationStructureBuildSizesInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR;
	vkGetAccelerationStructureBuildSizesKHR(
		GVulkanRHI->GetLogicalDevice(),
		VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
		&accelerationStructureSizeGeometryInfo,
		&maxPrimitiveCount,
		&accelerationStructureBuildSizesInfo);

	CreateAccelerationStructureBuffer(accelerationStructureBuildSizesInfo.accelerationStructureSize, mAccelerationStructureBuffer, mAccelerationStructureMemory);

	VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
	accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
	accelerationStructureCreateInfo.buffer = mAccelerationStructureBuffer;
	accelerationStructureCreateInfo.size = accelerationStructureBuildSizesInfo.accelerationStructureSize;
	accelerationStructureCreateInfo.type = mType;
	vkCreateAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), &accelerationStructureCreateInfo, nullptr, &mAccelerationStructure);

	VkBuffer scratchBuffer;
	VkDeviceMemory scratchMemory;
	VkDeviceAddress scratchAddress;
	CreateScratchBuffer(accelerationStructureBuildSizesInfo.buildScratchSize, scratchBuffer, scratchMemory, scratchAddress);

	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
	accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureBuildGeometryInfo.type = mType;
	accelerationStructureBuildGeometryInfo.flags = mFlags;
	accelerationStructureBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
	accelerationStructureBuildGeometryInfo.dstAccelerationStructure = mAccelerationStructure;
	accelerationStructureBuildGeometryInfo.geometryCount = 1;
	accelerationStructureBuildGeometryInfo.pGeometries = &mAccelerationStructureGeometry;
	accelerationStructureBuildGeometryInfo.scratchData.deviceAddress = scratchAddress;

	if (GVulkanRHI->GetAccelerationStructureFeatures().accelerationStructureHostCommands)
	{
		// Implementation supports building acceleration structure building on host
		vkBuildAccelerationStructuresKHR(
			GVulkanRHI->GetLogicalDevice(),
			VK_NULL_HANDLE,
			1,
			&accelerationStructureBuildGeometryInfo,
			accelerationStructureBuildRangeInfos.data());
	}
	else
	{
		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

		// Make sure the copy of the instance buffer are copied before triggering the acceleration structure build
		VkMemoryBarrier memoryBarrier{ VK_STRUCTURE_TYPE_MEMORY_BARRIER };
		memoryBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
		memoryBarrier.dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR;
		vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
			0, 1, &memoryBarrier, 0, nullptr, 0, nullptr);

		vkCmdBuildAccelerationStructuresKHR(
			commandBuffer,
			1,
			&accelerationStructureBuildGeometryInfo,
			accelerationStructureBuildRangeInfos.data());

		GRenderState->FlushSingleCommand(commandBuffer);
	}

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), scratchBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), scratchMemory, nullptr);

	// destroy instance buffer
	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mInstanceBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mInstanceMemory, nullptr);
}

void TopLevelAccelerationStructure::UpdateTlas()
{
	// All blas must had build already
	// Just rebuild the instance buffer for the instance transform
	// TODO : add or remove blas and instance

	// Rebuild the instance buffer
	// Maybe can update by write data
	CreateInstanceGeometries();

	uint32 maxPrimitiveCount = static_cast<uint32>(mBottomLevelAccelerationStructures.size());

	VkAccelerationStructureBuildRangeInfoKHR accelerationStructureBuildRangeInfo{};
	accelerationStructureBuildRangeInfo.primitiveCount = maxPrimitiveCount;
	accelerationStructureBuildRangeInfo.primitiveOffset = 0;
	accelerationStructureBuildRangeInfo.firstVertex = 0;
	accelerationStructureBuildRangeInfo.transformOffset = 0;

	std::array<VkAccelerationStructureBuildRangeInfoKHR*, 1> accelerationStructureBuildRangeInfos = { &accelerationStructureBuildRangeInfo };

	// Get Size Info
	VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
	accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
	accelerationStructureBuildGeometryInfo.type = mType;
	accelerationStructureBuildGeometryInfo.flags = mFlags;
	accelerationStructureBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR;
	accelerationStructureBuildGeometryInfo.geometryCount = 1;
	accelerationStructureBuildGeometryInfo.pGeometries = &mAccelerationStructureGeometry;
	accelerationStructureBuildGeometryInfo.dstAccelerationStructure = mAccelerationStructure;
	accelerationStructureBuildGeometryInfo.dstAccelerationStructure = mAccelerationStructure;

	VkAccelerationStructureBuildSizesInfoKHR accelerationStructureBuildSizesInfo{};
	accelerationStructureBuildSizesInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR;
	vkGetAccelerationStructureBuildSizesKHR(
		GVulkanRHI->GetLogicalDevice(),
		VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
		&accelerationStructureBuildGeometryInfo,
		&maxPrimitiveCount,
		&accelerationStructureBuildSizesInfo);

	VkBuffer scratchBuffer;
	VkDeviceMemory scratchMemory;
	VkDeviceAddress scratchAddress;
	CreateScratchBuffer(accelerationStructureBuildSizesInfo.buildScratchSize, scratchBuffer, scratchMemory, scratchAddress);

	accelerationStructureBuildGeometryInfo.scratchData.deviceAddress = scratchAddress;

	if (GVulkanRHI->GetAccelerationStructureFeatures().accelerationStructureHostCommands)
	{
		// Implementation supports building acceleration structure building on host
		vkBuildAccelerationStructuresKHR(
			GVulkanRHI->GetLogicalDevice(),
			VK_NULL_HANDLE,
			1,
			&accelerationStructureBuildGeometryInfo,
			accelerationStructureBuildRangeInfos.data());
	}
	else
	{
		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

		// Make sure the copy of the instance buffer are copied before triggering the acceleration structure build
		VkMemoryBarrier memoryBarrier{ VK_STRUCTURE_TYPE_MEMORY_BARRIER };
		memoryBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
		memoryBarrier.dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR;
		vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
			0, 1, &memoryBarrier, 0, nullptr, 0, nullptr);

		vkCmdBuildAccelerationStructuresKHR(
			commandBuffer,
			1,
			&accelerationStructureBuildGeometryInfo,
			accelerationStructureBuildRangeInfos.data());

		GRenderState->FlushSingleCommand(commandBuffer);
	}

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), scratchBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), scratchMemory, nullptr);

	// destroy instance buffer
	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mInstanceBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mInstanceMemory, nullptr);
}

void TopLevelAccelerationStructure::CompactionAndBuildBlas()
{
	std::vector<RefCountPtr<BottomLevelAccelerationStructure>> blases;

	for (auto& blas : mBottomLevelAccelerationStructures)
	{
		if (!blas->IsBuilded())
		{
			blases.push_back(blas);
		}
	}

	// No blas need to compaction
	if (blases.size() == 0)
		return;

	VkDeviceSize maxStructureSize = 0; // Memory size of all allocated BLAS
	VkDeviceSize maxScratchSize = 0; // Largest scratch size

	for (auto& blas : blases)
	{
		// Get Size Info
		VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureSizeGeometryInfo{};
		accelerationStructureSizeGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
		accelerationStructureSizeGeometryInfo.type = blas->GetType();
		accelerationStructureSizeGeometryInfo.flags = blas->GetFlags();
		accelerationStructureSizeGeometryInfo.geometryCount = static_cast<uint32>(blas->GetVkGeometries().size());
		accelerationStructureSizeGeometryInfo.pGeometries = blas->GetVkGeometries().data();

		vkGetAccelerationStructureBuildSizesKHR(
			GVulkanRHI->GetLogicalDevice(),
			VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
			&accelerationStructureSizeGeometryInfo,
			blas->GetMaxPrimitiveCounts().data(),
			&blas->GetSizeInfo());

		maxStructureSize += blas->GetSizeInfo().accelerationStructureSize;
		maxScratchSize += blas->GetSizeInfo().buildScratchSize;
	}

	VkBuffer scratchBuffer;
	VkDeviceMemory scratchMemory;
	VkDeviceAddress scratchAddress;
	CreateScratchBuffer(maxScratchSize, scratchBuffer, scratchMemory, scratchAddress);

	const uint32 blasCount = static_cast<uint32>(blases.size());

	// Allocate a query pool for storing the needed size for every BLAS compaction.
	VkQueryPool queryPool = VK_NULL_HANDLE;
	VkQueryPoolCreateInfo queryPoolCreateInfo{};
	queryPoolCreateInfo.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO;
	queryPoolCreateInfo.queryCount = blasCount;
	queryPoolCreateInfo.queryType = VK_QUERY_TYPE_ACCELERATION_STRUCTURE_COMPACTED_SIZE_KHR;
	vkCreateQueryPool(GVulkanRHI->GetLogicalDevice(), &queryPoolCreateInfo, nullptr, &queryPool);

	// Batching creation/compaction of BLAS to allow staying in restricted amount of memory
	VkDeviceSize batchSize = 0;
	VkDeviceSize batchLimitSize = 256000000;
	// Indices of the BLAS to create
	std::vector<uint32> indices;
	for (uint32 i = 0; i < blasCount; ++i)
	{
		indices.push_back(i);
		batchSize += blases[i]->GetSizeInfo().accelerationStructureSize;
		// Over the limit or last BLAS element
		if (batchSize >= batchLimitSize || i == blasCount - 1)
		{
			VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

			// For querying the compaction size
			if (queryPool)
			{
				vkResetQueryPool(GVulkanRHI->GetLogicalDevice(), queryPool, 0, static_cast<uint32>(indices.size()));
			}

			uint32 queryCount = 0;

			for (auto j : indices)
			{
				CreateAccelerationStructureBuffer(blases[j]->GetSizeInfo().accelerationStructureSize, blases[j]->GetAccelerationStructureBuffer(), blases[j]->GetAccelerationStructureMemory());

				VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
				accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
				accelerationStructureCreateInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
				accelerationStructureCreateInfo.buffer = blases[j]->GetAccelerationStructureBuffer();
				accelerationStructureCreateInfo.size = blases[j]->GetSizeInfo().accelerationStructureSize;
				vkCreateAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), &accelerationStructureCreateInfo, nullptr, &blases[j]->GetAccelerationStructure());

				VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{};
				accelerationStructureBuildGeometryInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
				accelerationStructureBuildGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
				accelerationStructureBuildGeometryInfo.flags = blases[j]->GetFlags();
				accelerationStructureBuildGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
				accelerationStructureBuildGeometryInfo.dstAccelerationStructure = blases[j]->GetAccelerationStructure();
				accelerationStructureBuildGeometryInfo.geometryCount = static_cast<uint32>(blases[j]->GetVkGeometries().size());
				accelerationStructureBuildGeometryInfo.pGeometries = blases[j]->GetVkGeometries().data();
				accelerationStructureBuildGeometryInfo.scratchData.deviceAddress = scratchAddress;

				// Building the bottom-level-acceleration-structure
				vkCmdBuildAccelerationStructuresKHR(
					commandBuffer,
					1,
					&accelerationStructureBuildGeometryInfo,
					blases[j]->GetRangeInfos().data());

				// Since the scratch buffer is reused across builds, we need a barrier to ensure one build
				// is finished before starting the next one.
				VkMemoryBarrier memoryBarrier{ VK_STRUCTURE_TYPE_MEMORY_BARRIER };
				memoryBarrier.srcAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR;
				memoryBarrier.dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR;
				vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
					VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR, 0, 1, &memoryBarrier, 0, nullptr, 0, nullptr);

				// Add a query to find the 'real' amount of memory needed, use for compaction
				if (queryPool != VK_NULL_HANDLE)
				{
					vkCmdWriteAccelerationStructuresPropertiesKHR(commandBuffer, 1, &blases[j]->GetAccelerationStructure(),
						VK_QUERY_TYPE_ACCELERATION_STRUCTURE_COMPACTED_SIZE_KHR, queryPool, queryCount++);
				}
			}

			GRenderState->FlushSingleCommand(commandBuffer);
		}

		if (queryPool)
		{
			VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

			uint32 queryCount = 0;
			// Get the compacted size result back
			std::vector<VkDeviceSize> compactSizes(static_cast<uint32>(indices.size()));
			vkGetQueryPoolResults(GVulkanRHI->GetLogicalDevice(), queryPool, 0, static_cast<uint32>(compactSizes.size()), compactSizes.size() * sizeof(VkDeviceSize), compactSizes.data(), sizeof(VkDeviceSize), VK_QUERY_RESULT_WAIT_BIT);

			for (auto j : indices)
			{
				// previous AS to destroy
				blases[j]->SaveClearStructure();

				// new reduced size
				blases[j]->GetSizeInfo().accelerationStructureSize = compactSizes[queryCount++];

				// Creating a compact version of the AS
				CreateAccelerationStructureBuffer(blases[j]->GetSizeInfo().accelerationStructureSize, blases[j]->GetAccelerationStructureBuffer(), blases[j]->GetAccelerationStructureMemory());

				VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfo{};
				accelerationStructureCreateInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
				accelerationStructureCreateInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
				accelerationStructureCreateInfo.buffer = blases[j]->GetAccelerationStructureBuffer();
				accelerationStructureCreateInfo.size = blases[j]->GetSizeInfo().accelerationStructureSize;
				vkCreateAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), &accelerationStructureCreateInfo, nullptr, &blases[j]->GetAccelerationStructure());

				// Copy the original BLAS to a compact version
				VkCopyAccelerationStructureInfoKHR copyAccelerationStructureInfo{};
				copyAccelerationStructureInfo.sType = VK_STRUCTURE_TYPE_COPY_ACCELERATION_STRUCTURE_INFO_KHR;
				copyAccelerationStructureInfo.src = blases[j]->GetClearStructure();
				copyAccelerationStructureInfo.dst = blases[j]->GetAccelerationStructure();
				copyAccelerationStructureInfo.mode = VK_COPY_ACCELERATION_STRUCTURE_MODE_COMPACT_KHR;
				vkCmdCopyAccelerationStructureKHR(commandBuffer, &copyAccelerationStructureInfo);
			}

			GRenderState->FlushSingleCommand(commandBuffer);

			// Destroy the non-compacted blas and get the instance device
			for (auto j : indices)
			{
				blases[j]->DeleteClearStructure();
				blases[j]->GetInstanceAddress();
			}
		}

		// Reset
		batchSize = 0;
		indices.clear();
	}

	vkDestroyQueryPool(GVulkanRHI->GetLogicalDevice(), queryPool, nullptr);

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), scratchBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), scratchMemory, nullptr);
}

void TopLevelAccelerationStructure::CreateInstanceGeometries()
{
	std::vector<VkAccelerationStructureInstanceKHR> accelerationStructureInstances;
	for (auto& blas : mBottomLevelAccelerationStructures)
		accelerationStructureInstances.push_back(blas->GetAccelerationStructureInstance());

	const VkDeviceSize bufferSize = static_cast<VkDeviceSize>(accelerationStructureInstances.size() * sizeof(VkAccelerationStructureInstanceKHR));

	VkBufferCreateInfo bufferCreateInfo{};
	bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferCreateInfo.size = bufferSize;
	bufferCreateInfo.usage = VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR;
	bufferCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	VK_CHECK_RESULT(vkCreateBuffer(GVulkanRHI->GetLogicalDevice(), &bufferCreateInfo, nullptr, &mInstanceBuffer));

	VkMemoryRequirements memoryRequirements;
	vkGetBufferMemoryRequirements(GVulkanRHI->GetLogicalDevice(), mInstanceBuffer, &memoryRequirements);

	VkMemoryAllocateInfo memoryAllocateInfo{};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.allocationSize = memoryRequirements.size;
	memoryAllocateInfo.memoryTypeIndex = GVulkanRHI->FindMemoryType(memoryRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

	// If usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, need to add a VkMemoryAllocateFlagsInfo to pNext of VkMemoryAllocateInfo
	VkMemoryAllocateFlagsInfo memoryAllocateFlagsInfo{};
	memoryAllocateFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
	memoryAllocateFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;
	memoryAllocateInfo.pNext = &memoryAllocateFlagsInfo;

	VK_CHECK_RESULT(vkAllocateMemory(GVulkanRHI->GetLogicalDevice(), &memoryAllocateInfo, nullptr, &mInstanceMemory));
	vkBindBufferMemory(GVulkanRHI->GetLogicalDevice(), mInstanceBuffer, mInstanceMemory, 0);

	void* data;
	VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), mInstanceMemory, 0, bufferSize, 0, &data));
	memcpy(data, accelerationStructureInstances.data(), bufferSize);
	vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), mInstanceMemory);

	VkBufferDeviceAddressInfoKHR bufferDeviceAddressInfo{};
	bufferDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	bufferDeviceAddressInfo.buffer = mInstanceBuffer;
	mAccelerationStructureGeometry.geometry.instances.data.deviceAddress = vkGetBufferDeviceAddressKHR(GVulkanRHI->GetLogicalDevice(), &bufferDeviceAddressInfo);
	assert(mAccelerationStructureGeometry.geometry.instances.data.deviceAddress != 0);
}

